import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import Lasso
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

# ==================== 初始化设置 ====================
# 设置中文字体和数学符号显示
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 使用楷体显示中文
plt.rcParams['mathtext.fontset'] = 'stix'    # 数学公式字体
plt.rcParams['axes.unicode_minus'] = False   # 解决负号显示问题

# ==================== 数据准备阶段 ====================
# 设置随机种子保证结果可复现
np.random.seed(42)

# 生成特征数据X：100个样本，1个特征，取值范围[0,2)
X = 2 * np.random.rand(100, 1)

# 生成标签数据：真实模型为 y = 4 + 3X + 噪声
# 4是截距项(bias term)，3是权重系数(weight)
# 噪声服从标准正态分布 N(0,1)
y = 4 + 3 * X + np.random.randn(100, 1)

# 划分训练集和测试集 (80%训练，20%测试)
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=42
)

# ==================== Lasso回归模型 ====================
# 创建Lasso回归模型
# alpha: 正则化系数，控制模型复杂度 (默认1.0)
# max_iter: 最大迭代次数 (默认1000)
# tol: 优化算法的容忍度 (默认1e-4)
# 注意：对于小数据集可以设置较大的max_iter和较小的tol
lasso_reg = Lasso(
    alpha=0.1,          # 正则化强度
    max_iter=10000,     # 增加迭代次数确保收敛
    tol=1e-6,           # 更严格的收敛标准
    random_state=42     # 随机种子
)

# 训练模型
lasso_reg.fit(X_train, y_train)

# ==================== 模型评估 ====================
# 训练集预测
y_train_pred = lasso_reg.predict(X_train)
train_mse = mean_squared_error(y_train, y_train_pred)

# 测试集预测
y_test_pred = lasso_reg.predict(X_test)
test_mse = mean_squared_error(y_test, y_test_pred)

# ==================== 结果输出 ====================
print("===== Lasso回归结果 =====")
print(f"截距项(w0): {lasso_reg.intercept_[0]:.4f}")  # 添加[0]提取数组中的值
print(f"系数(w1): {lasso_reg.coef_[0]:.4f}")
print(f"训练集MSE: {train_mse:.4f}")
print(f"测试集MSE: {test_mse:.4f}")
print(f"使用的特征数: {np.sum(lasso_reg.coef_ != 0)}")  # Lasso的特征选择能力

# ==================== 可视化 ====================
plt.figure(figsize=(12, 6))

# 绘制训练数据
plt.scatter(
    X_train, y_train,
    color='blue',
    label='训练数据',
    alpha=0.6
)

# 绘制测试数据
plt.scatter(
    X_test, y_test,
    color='green',
    label='测试数据',
    alpha=0.6
)

# 绘制拟合线
X_plot = np.linspace(0, 2, 100).reshape(-1, 1)
y_plot = lasso_reg.predict(X_plot)
plt.plot(
    X_plot, y_plot,
    color='red',
    linewidth=2,
    label='Lasso回归拟合线'
)

# 添加图注和标题
plt.title('Lasso回归示例 (alpha=0.1)', fontsize=14)
plt.xlabel('X特征值', fontsize=12)
plt.ylabel('y标签值', fontsize=12)
plt.legend(fontsize=10)
plt.grid(True)

# 显示图形
plt.show()
